from utils.read_config import read_config, read_matrix_config
from utils.util import *


class Data():

    def __init__(self, config_dir=None):

        self.var_dict = None
        self.func_dict = None
        self.const_dict = None
        self.tmp_dict = None
        self.description = None
        self.add_dict = None
        self.sub_dict = None
        self.multi_dict = None
        self.multiScalar_dict = None
        self.transpose_dict = None
        self.inversion_dict = None
        self.equal_dict = None
        self.vector_norm_dict = None
        self.global_vars = None
        self.var_cnt = 0
        self.exception = None
        self.llm_mode = False

        # load information from configs
        self.var_dict, self.func_dict, self.const_dict, \
            self.tmp_dict, self.description = read_config(config_dir)
        self.var_dict = dict(sorted(self.var_dict.items(), key=lambda x: x[0], reverse=True))
        self.const_dict = dict(sorted(self.const_dict.items(), key=lambda x: x[0], reverse=True))
        self.tmp_dict = dict(sorted(self.tmp_dict.items(), key=lambda x: x[0], reverse=True))
        self.add_dict, self.sub_dict, self.multi_dict, self.multiScalar_dict, self.transpose_dict, \
            self.inversion_dict, self.equal_dict, self.vector_norm_dict \
            = read_matrix_config('./configs/dicts/matrix_config.json')
        self.global_vars = {}
        for d in self.var_dict.values():
            self.global_vars[d['name']] = {'type': d['type'], 'dim': d['dim']}

    def get_tmp_var(self):

        cnt = self.var_cnt
        name = 'tmp' + str(cnt)
        while name in self.tmp_dict.keys():
            cnt += 1
            name = 'tmp' + str(cnt)
        self.var_cnt += 1
        return name

    def get_recommend_name(self, var_name):

        var_dicts = {**self.global_vars, **self.tmp_dict, **self.const_dict}
        new_name = check_name_validity(var_name)
        if new_name == '' or new_name in var_dicts.keys():
            new_name = self.get_tmp_var()
        return new_name

    def get_tmp_arrays(self):

        # 过滤一下对数组的引用，必须是形如t[x][y]且维度是1*1
        result = {}  # 存储过滤后的dict
        for v in self.tmp_dict.keys():
            if v in result.keys():
                continue
            if self.tmp_dict[v]['dim'] != '1*1':
                result[v] = self.tmp_dict[v]
                continue
            var = v
            if 'name' in self.tmp_dict[v].keys():
                var = self.tmp_dict[v]['name']
            if not is_array(var):
                result[v] = self.tmp_dict[v]
                continue
            # 是数组的引用
            array_name = drop_square_brackets(var)
            if array_name in self.tmp_dict.keys():
                result[array_name] = self.tmp_dict[array_name]
                continue

            def is_same_array(a1: str, a2: str, a2_info: dict) -> bool:
                if a2_info['dim'] != '1*1':
                    return False
                if 'name' in a2_info.keys():
                    a2 = a2_info['name']
                if drop_square_brackets(a1) == drop_square_brackets(a2):
                    return True
                return False

            cur_array = {k: v_ for (k, v_) in self.tmp_dict.items() if is_same_array(var, k, v_)}

            def get_indexs(s: str):
                # 获取当前数组每一维度的索引值
                results = []
                while s[-1] == ']':
                    l = bracket_match(s, len(s) - 1)
                    idx = s[l + 1:-1]
                    if is_int_num(idx):
                        results += [int(idx)]
                    else:
                        results += [2]
                    s = s[:l]
                return results

            final_dims = []
            for k, v_ in cur_array.items():
                if 'name' in v_.keys():
                    k = v_['name']
                cur_idx = get_indexs(k)
                if final_dims == []:
                    final_dims = cur_idx
                else:
                    if len(final_dims) != len(cur_idx):
                        # 暂时不处理错误
                        # erro_print(traceback.extract_stack(), f"Array use wrong, array name is: {array_name}")
                        result[var] = self.tmp_dict[var]
                        break
                    else:
                        for i in range(len(final_dims)):
                            final_dims[i] = max(final_dims[i], cur_idx[i])
            if final_dims != []:
                final_dims = [str(i + 1) for i in final_dims]
                result[array_name] = {'type': self.tmp_dict[v]['type'], 'dim': '*'.join(final_dims)}
        self.tmp_dict = result

    def is_struct(self, var: str) -> bool:

        # warn 这个函数非常不严谨
        var_dicts = {**self.global_vars, **self.tmp_dict, **self.const_dict}
        if var.split('[')[0] in var_dicts.keys():
            return var.split('[')[0]
        if var.split('.')[0] in var_dicts.keys():
            return var.split('.')[0]
        return None

    def get_args(self, var_list, type_list):

        res = []
        if len(var_list) != len(type_list):
            return None
        var_dicts = {**self.global_vars, **self.tmp_dict, **self.const_dict}
        for i in range(len(var_list)):
            var = var_list[i]
            ty, dim = '', ''
            if var in var_dicts.keys():
                ty, dim = var_dicts[var]['type'], var_dicts[var]['dim']
            else:
                for v in var_dicts.keys():
                    if 'name' in var_dicts[v].keys() and var_dicts[v]['name'] == var:
                        ty, dim = var_dicts[var]['type'], var_dicts[var]['dim']
                # 常数放过
                if is_num(var) and '*' not in type_list[i]:
                    if '*' not in type_list[i]:
                        res.append(var)
                        continue
                # 运算放过
                if '*' in var or '/' in var:
                    if '*' not in type_list[i]:
                        res.append(var)
                        continue
                if ty == '':
                    return None
            if dim == '1*1' and ty != type_list[i] and ty + '*' != type_list[i]:
                return None
            elif dim == '1*1' and ty == type_list[i]:
                res.append(var)
            elif dim == '1*1' and ty + '*' == type_list[i]:
                res.append('&' + var)
            elif ty + '*' != type_list[i]:
                return None
            else:
                if '1' in dim.split('*'):
                    res.append('&' + var + '[0]')
                else:
                    res.append('&' + var + '[0][0]')
        return res
